#include "coarsen.hpp"
#include "aggregate/aggregate.hpp"
#include "parcsr_matop.hpp"
#include "prolong.hpp"

template<typename idx_t, typename setup_t, int dof>
void Coarsen(const idx_t coarsen_type, const par_CSRMatrix<idx_t, setup_t, setup_t, dof> * par_fine_mat,
    par_CSRMatrix<idx_t, setup_t, setup_t, 1> * & R_mat, par_CSRMatrix<idx_t, setup_t, setup_t, dof> * & par_coar_mat,
    par_CSRMatrix<idx_t, setup_t, setup_t, 1> * & P_mat, const CoarseConfig & config)
{
    idx_t * C_rows_partits = nullptr;// 粗网格上的进程划分
    const MPI_Comm comm = par_fine_mat->comm;
    int my_pid; MPI_Comm_rank(comm, & my_pid);
    int num_procs; MPI_Comm_size(comm, & num_procs);

    double t_aggre = MPI_Wtime();
    if (coarsen_type == 0) {
        C_rows_partits = naive_pairwise(par_fine_mat, R_mat, P_mat);
    } else if (coarsen_type == 1) {
        C_rows_partits = coupled_pairwise(par_fine_mat, R_mat, P_mat, (setup_t) config.theta);
    } else if (coarsen_type == 2) {
        C_rows_partits = modified_pairwise(par_fine_mat, R_mat, P_mat, config.iter);
    } else if (coarsen_type == 10) {
        C_rows_partits = naive_neighbor(par_fine_mat, R_mat, P_mat);
    } else if (coarsen_type == 12) {
        C_rows_partits = modified_neighbor(par_fine_mat, R_mat, P_mat);
    }
    else MPI_Abort(comm, -1112);

    t_aggre = MPI_Wtime() - t_aggre;
    if (my_pid == 0) printf("Aggre type %d cost %f config iter %d theta %.2f\n", coarsen_type, t_aggre, config.iter, config.theta);
    // 如有需要，计算光滑后的插值矩阵
    // par_CSRMatrix<idx_t, setup_t, setup_t, dof> * smth_P_mat = nullptr, * smth_R_mat = nullptr;
    // if (config.smoothed_weight > 0.0) {
    //     par_CSRMatrix<idx_t, setup_t, setup_t, dof> * jac_mat = Jacobi_A((setup_t) config.smoothed_weight, * par_fine_mat);
    //     smth_P_mat = new par_CSRMatrix<idx_t, setup_t, setup_t, dof>(comm, P_mat->rows_partition, P_mat->cols_partition);
    //     smth_R_mat = new par_CSRMatrix<idx_t, setup_t, setup_t, dof>(comm);
    //     par_CSRMatrixMultiply(* jac_mat, * P_mat, * smth_P_mat);
    //     par_CSRMatrixTranspose(*smth_P_mat, *smth_R_mat, true);
    //     delete jac_mat;
    // }

    // 计算粗网格矩阵
    par_coar_mat = new par_CSRMatrix<idx_t, setup_t, setup_t, dof>(comm, C_rows_partits, C_rows_partits);
    par_coar_mat->owns_row_partits = true;

    {// 计算RAP
        par_CSRMatrix<idx_t, setup_t, setup_t, dof> tmp_1(comm, par_fine_mat->rows_partition, P_mat->cols_partition);// A*P
        // if (smth_P_mat != nullptr) { assert(config.smoothed_weight > 0.0);// 光滑聚合
        //     par_CSRMatrixMultiply(* par_fine_mat, * smth_P_mat, tmp_1);
        //     par_CSRMatrixMultiply(* smth_R_mat, tmp_1, * par_coar_mat);
        // } else {// 非光滑聚合
            par_CSRMatrixMultiply(* par_fine_mat, * P_mat, tmp_1);// 要以AF的通信包收集PExt
            par_CSRMatrixMultiply(* R_mat, tmp_1, * par_coar_mat);// 要以R的通信包收集tmp_1Ext
        // }
#ifdef DEBUG
        for (int p = 0; p < num_procs; p++) {
            if (my_pid == p) {
                printf("Proc %d tmp_1 nrows %d ncols %d diag %d %d %d offd %d %d %d\n",
                    my_pid, tmp_1.glb_nrows, tmp_1.glb_ncols,   tmp_1.diag.nrows, tmp_1.diag.ncols, tmp_1.diag.nnz,
                                                                tmp_1.offd.nrows, tmp_1.offd.ncols, tmp_1.offd.nnz);
                for (idx_t i = 0; i < tmp_1.diag.nrows; i++) {
                    printf("  %d : ", tmp_1.beg_row + i);
                    for (idx_t p = tmp_1.diag.row_ptr[i]; p < tmp_1.diag.row_ptr[i + 1]; p++)
                        printf("%d ", tmp_1.beg_col + tmp_1.diag.col_idx[p]);
                    printf(" | ");
                    for (idx_t p = tmp_1.offd.row_ptr[i]; p < tmp_1.offd.row_ptr[i + 1]; p++)
                        printf("%d ", tmp_1.col_map_offd[tmp_1.offd.col_idx[p]]);
                    printf("\n");
                }
            }
            MPI_Barrier(comm);
        }
#endif

#ifdef DEBUG
        for (int p = 0; p < num_procs; p++) {
            if (my_pid == p) {
                printf("Proc %d AC nrows %d ncols %d diag %d %d %d offd %d %d %d\n",
                    my_pid, par_coar_mat->glb_nrows, par_coar_mat->glb_ncols,
                    par_coar_mat->diag.nrows, par_coar_mat->diag.ncols, par_coar_mat->diag.nnz,
                    par_coar_mat->offd.nrows, par_coar_mat->offd.ncols, par_coar_mat->offd.nnz);
                for (idx_t i = 0; i < par_coar_mat->diag.nrows; i++) {
                    printf("  %d : ", par_coar_mat->beg_row + i);
                    for (idx_t p = par_coar_mat->diag.row_ptr[i]; p < par_coar_mat->diag.row_ptr[i + 1]; p++)
                        printf("(%d %.2f) ", par_coar_mat->beg_col + par_coar_mat->diag.col_idx[p], par_coar_mat->diag.vals[p]);
                    printf(" | ");
                    for (idx_t p = par_coar_mat->offd.row_ptr[i]; p < par_coar_mat->offd.row_ptr[i + 1]; p++)
                        printf("(%d %.2f) ", par_coar_mat->col_map_offd[par_coar_mat->offd.col_idx[p]], par_coar_mat->offd.vals[p]);
                    printf("\n");
                }
            }
            MPI_Barrier(comm);
        }
#endif
    }

    if (iso_wells.empty() == false)
    {// 更新井在粗网格的序号
        int send_buf [coar_wells.size()];
        int my_num_wells = 0, glb_num_wells = iso_wells.size();
        for (auto it = coar_wells.begin(); it != coar_wells.end(); it ++) {
            send_buf[my_num_wells ++] = *it + R_mat->beg_row;
        }
        assert(my_num_wells == (int) coar_wells.size());
        int coar_well_glb_ids [glb_num_wells];
        int recv_cnts [num_procs], displs [num_procs];
        for (int p = 0; p < num_procs; p++) {
            recv_cnts[p] = 0;
            for (auto it = iso_wells.begin(); it != iso_wells.end(); it++) {
                if (par_fine_mat->rows_partition[p] <= *it && *it < par_fine_mat->rows_partition[p + 1]) recv_cnts[p] ++;
            }
            displs[p] = (p == 0) ? 0 : (displs[p - 1] + recv_cnts[p - 1]);
        }
        // printf("Proc %d recv %d %d displs %d %d\n", my_pid, recv_cnts[0], recv_cnts[1], displs[0], displs[1]);
        // printf("Proc %d num %d \n", my_pid, my_num_wells);
        MPI_Allgatherv(send_buf, my_num_wells, MPI_INT, coar_well_glb_ids, recv_cnts, displs, MPI_INT, comm);

        iso_wells.clear();
        if (glb_num_wells > 0 && my_pid == 0) printf("Coar Well ID: ");
        for (int i = 0; i < glb_num_wells; i++) {
            const idx_t well_id = coar_well_glb_ids[i];
            if (glb_num_wells > 0 && my_pid == 0) printf("%d ", well_id);
            if (par_coar_mat->beg_row <= well_id && well_id < par_coar_mat->end_row) {
                const idx_t loc_ci = well_id - par_coar_mat->beg_row;
                assert(R_mat->diag.row_ptr[loc_ci] + 1 == R_mat->diag.row_ptr[loc_ci + 1]);
                assert(R_mat->diag.vals[R_mat->diag.row_ptr[loc_ci]] == 1.0);
            }
            iso_wells.emplace(well_id);
        }
        if (glb_num_wells > 0 && my_pid == 0) { printf("\n"); }
    }

    // delete smth_P_mat; delete smth_R_mat;
}
template void Coarsen(const int, const par_CSRMatrix_I32F64F64_DOF1 *, par_CSRMatrix_I32F64F64_DOF1 * &, par_CSRMatrix_I32F64F64_DOF1 * &, par_CSRMatrix_I32F64F64_DOF1 * &, const CoarseConfig & );
template void Coarsen(const int, const par_CSRMatrix_I32F64F64_DOF2 *, par_CSRMatrix_I32F64F64_DOF1 * &, par_CSRMatrix_I32F64F64_DOF2 * &, par_CSRMatrix_I32F64F64_DOF1 * &, const CoarseConfig & );
template void Coarsen(const int, const par_CSRMatrix_I32F64F64_DOF3 *, par_CSRMatrix_I32F64F64_DOF1 * &, par_CSRMatrix_I32F64F64_DOF3 * &, par_CSRMatrix_I32F64F64_DOF1 * &, const CoarseConfig & );
template void Coarsen(const int, const par_CSRMatrix_I32F64F64_DOF4 *, par_CSRMatrix_I32F64F64_DOF1 * &, par_CSRMatrix_I32F64F64_DOF4 * &, par_CSRMatrix_I32F64F64_DOF1 * &, const CoarseConfig & );

template void Coarsen(const int, const par_CSRMatrix_I32F32F32_DOF1 *, par_CSRMatrix_I32F32F32_DOF1 * &, par_CSRMatrix_I32F32F32_DOF1 * &, par_CSRMatrix_I32F32F32_DOF1 * &, const CoarseConfig & );
template void Coarsen(const int, const par_CSRMatrix_I32F32F32_DOF2 *, par_CSRMatrix_I32F32F32_DOF1 * &, par_CSRMatrix_I32F32F32_DOF2 * &, par_CSRMatrix_I32F32F32_DOF1 * &, const CoarseConfig & );
template void Coarsen(const int, const par_CSRMatrix_I32F32F32_DOF3 *, par_CSRMatrix_I32F32F32_DOF1 * &, par_CSRMatrix_I32F32F32_DOF3 * &, par_CSRMatrix_I32F32F32_DOF1 * &, const CoarseConfig & );
template void Coarsen(const int, const par_CSRMatrix_I32F32F32_DOF4 *, par_CSRMatrix_I32F32F32_DOF1 * &, par_CSRMatrix_I32F32F32_DOF4 * &, par_CSRMatrix_I32F32F32_DOF1 * &, const CoarseConfig & );

